# xray_image_processor.py

import numpy as np
import matplotlib.pyplot as plt
from xray_scan import XRayScan
from xray_scan_collection import XRayScanCollection


class XRayImageProcessor:
    def __init__(self, xray_collection):
        """
        初始化XRayImageProcessor对象。

        :param xray_collection: XRayScanCollection 对象
        """
        self.xray_collection = xray_collection

    def plot_polar_image(self):
        """
        用极坐标绘制传感器数据图像。
        """
        angles = []
        data = []
        max_relative_length = 0

        for scan in self.xray_collection.scans:
            angles.append(np.deg2rad(scan.angle_deg))
            data.append(scan.sensor_data)
            max_relative_length = max(max_relative_length, scan.get_relative_length())

        angles = np.array(angles)
        data = np.array(data)

        # 创建极坐标图
        fig, ax = plt.subplots(subplot_kw={'projection': 'polar'})

        # 计算相对半径
        radii = np.linspace(-max_relative_length, max_relative_length, len(data[0]))

        for i in range(len(data)):
            for j in range(len(data[i])):
                ax.scatter(angles[i], radii[j], c=data[i][j], cmap='viridis', s=100, edgecolor='k')

        ax.set_theta_zero_location('N')
        ax.set_theta_direction(-1)
        plt.title('Polar Plot of XRayScanCollection')
        plt.colorbar(ax.scatter(angles, radii, c=data.flatten(), cmap='viridis', s=100, edgecolor='k'))
        plt.show()


def main():
    # 创建几个XRayScan对象
    scan1 = XRayScan(0, [0, 1, 0, 3, 0, 2, 0, 1, 0])
    scan2 = XRayScan(30, [0, 1, 0, 3, 0, 2, 0, 1, 0])
    scan3 = XRayScan(60, [0, 1, 0, 3, 0, 2, 0, 1, 0])
    scan4 = XRayScan(90, [0, 1, 0, 3, 0, 2, 0, 1, 0])

    # 创建XRayScanCollection对象并添加扫描
    collection = XRayScanCollection()
    collection.add_scan(scan1)
    collection.add_scan(scan2)
    collection.add_scan(scan3)
    collection.add_scan(scan4)

    # 创建XRayImageProcessor对象
    processor = XRayImageProcessor(collection)

    # 绘制极坐标图像
    processor.plot_polar_image()


if __name__ == "__main__":
    main()